import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import rcParams

# 设置字体为 SimHei（黑体），避免中文乱码
rcParams['font.sans-serif'] = ['SimHei']  # 设置默认字体
rcParams['axes.unicode_minus'] = False    # 解决负号显示问题

# 加载数据
def load_data(file_path):
    sheets = ['材料1', '材料2', '材料3', '材料4']
    df_list = []
    for sheet in sheets:
        df = pd.read_excel(file_path, sheet_name=sheet)
        df['Material'] = sheet
        df_list.append(df)
    full_df = pd.concat(df_list, ignore_index=True)
    return full_df

# 绘制每种材料的不同波形图
def plot_waveforms_by_material_and_waveform(df):
    # 将励磁波形映射为数值编码
    waveform_mapping = {'正弦波': 0, '三角波': 1, '梯形波': 2}
    df['励磁波形编码'] = df['励磁波形'].map(waveform_mapping)

    # 定义每种波形的颜色和标题
    waveforms = {
        0: ('正弦波', 'blue'),
        1: ('三角波', 'red'),
        2: ('梯形波', 'green')
    }

    # 获取所有材料名称
    materials = df['Material'].unique()

    # 创建 4 个子图，每个材料对应一个子图，每个材料包含 3 种波形的图像
    fig, axs = plt.subplots(4, 3, figsize=(18, 16))  # 4行3列，每个材料一行，每种波形一列

    for i, material in enumerate(materials):
        # 筛选特定材料的数据
        subset = df[df['Material'] == material]

        if subset.empty:
            print(f"没有找到 {material} 的数据")
            continue

        # 为每种励磁波形绘制数据
        for j, (code, (waveform, color)) in enumerate(waveforms.items()):
            # 筛选特定波形的数据
            waveform_data = subset[subset['励磁波形编码'] == code]
            if waveform_data.empty:
                print(f"材料 {material} 没有 {waveform} 的数据")
                continue

            # 绘制数据
            for idx, row in waveform_data.iterrows():
                magnetic_flux_density = row.iloc[4:1028].values  # 第5到1029列是磁通密度数据
                time_points = list(range(1, 1025))
                axs[i, j].plot(time_points, magnetic_flux_density, color=color, alpha=0.5)

            # 设置每个小图的标题和标签
            axs[i, j].set_title(f'{material} - {waveform}')
            axs[i, j].set_xlabel('时间点')
            axs[i, j].set_ylabel('磁通密度 (T)')
            axs[i, j].grid(True)

    plt.tight_layout()
    plt.show()

# 主程序
file_path = '附件一（训练集）.xlsx'
df = load_data(file_path)

# 绘制四种材料的不同波形图
plot_waveforms_by_material_and_waveform(df)
